from __future__ import print_function
import os
import torch
import torch.nn as nn
import torchvision.utils as vutils
from torch.autograd import Variable
from PIL import Image
import numpy as np

from models.u_net import UNet
from models.seg_net import Segnet
from models.fcn import FCN8s, VGGNet
from models.multi_capsule import CapsuleNet, CapsuleLoss

import torchvision.transforms as transforms
from torchvision.transforms import ToPILImage

class Transformer(object):
    def __init__(self, size, interpolation=Image.BILINEAR):
        self.size = size
        self.interpolation = interpolation
        self.toTensor = transforms.ToTensor()

    def __call__(self, img_):
        img_ = img_.resize(self.size, self.interpolation)
        img_ = self.toTensor(img_)
        img_.sub_(0.5).div_(0.5)
        return img_

#####################################################################################
# ToDo:init the model
#####################################################################################
model = UNet(3, 1)
# model = Segnet(3, 1)
# model = FCN8s(pretrained_net=VGGNet(pretrained=False),n_class=1)
# model = CapsuleNet(num_parts=5)

#####################################################################################
# ToDo:load the model's parameters
#####################################################################################
model_path = './checkpoint/Unet/model/netG_final.pth'
# model_path = './checkpoint/Segnet/model/netG_final.pth'
# model_path = './checkpoint/FCN/model/netG_final.pth'
# model_path = './checkpoint/MultiCapsule/model/netG_final.pth'

model.load_state_dict(torch.load(model_path))

test_image_path = './data/val/src/202105061613530082080_7_5.png'
test_image = Image.open(test_image_path).convert('RGB')
print('Operating...')
transformer = Transformer((256, 256))
img = transformer(test_image)
# 增加batch维度
img = img.unsqueeze(0)
img = Variable(img)

#####################################################################################
# ToDo:注意分辨模型是否是胶囊网络
#####################################################################################
label_image = model(img)
# part_map, label_image = model(img)

label_image = label_image.squeeze(0)

# 阈值分割
Threshold=torch.tensor(0.6)
label_image[label_image < Threshold] = torch.tensor(0.)
label_image[label_image >= Threshold] = torch.tensor(1.)

show = ToPILImage()
a = show(label_image)

#####################################################################################
# ToDo:注意分辨网络
#####################################################################################
a.save('./paper_result/unet.png')
# a.save('./paper_result/segnet.png')
# a.save('./paper_result/fcn.png')
# a.save('./paper_result/capsule.png')